from typing import Tuple
from .ID生成 import 循环ID生成器
from .UDP协议 import UDP协议
from .UDP消息 import UDP消息, UDP消息类型
from .异常 import 超时, UDP链接已关闭, UDP会话已关闭
from .通用 import UDP服务器信息
from .UDP心跳 import UDP心跳
import asyncio
import time


class UDP会话:
    def __init__(self, 服务器, 协议: UDP协议, 会话ID: int, 地址: Tuple[str, int] | None = None) -> None:
        self.服务器 = 服务器
        self.协议 = 协议
        self.会话ID = 会话ID
        self.地址 = 地址
        self.消息ID生成器 = 循环ID生成器()
        self.收到的消息队列 = asyncio.Queue()
        self.需要的消息ID生成器 = 循环ID生成器()
        self.需要的消息ID = self.需要的消息ID生成器()
        self.发送消息队列 = asyncio.Queue()
        self.待确认消息: UDP消息 | None = None
        self.消息确认futre = None
        self.loop = asyncio.get_running_loop()
        self.会话关闭 = False
        self.重传任务 = self.loop.create_task(self.消息重发任务())
        self.心跳 = UDP心跳(self.会话ID, 地址)
        self.心跳任务 = self.loop.create_task(
            self.心跳(self.运行状态检测, self.发送无需确认是否收到消息, self.关闭))
        self.处理任务: asyncio.Task | None = None

    def 关闭(self):
        self.会话关闭 = True
        self.重传任务.cancel()
        self.心跳任务.cancel()
        if self.处理任务 is not None:
            self.处理任务.cancel()
        self.服务器.关闭会话(self.会话ID)

    def 运行状态检测(self):
        if self.会话关闭:
            raise UDP会话已关闭(self.会话ID)
        if self.协议.链接关闭:
            raise UDP链接已关闭()

    async def 收消息(self, 消息: UDP消息):
        self.运行状态检测()
        self.心跳.刷新接收()
        if 消息.消息类型 == UDP消息类型.数据:
            if self.需要的消息ID == 消息.消息ID:
                self.需要的消息ID = self.需要的消息ID生成器()
                self.收到的消息队列.put_nowait(消息)
                self.发送无需确认是否收到消息(消息.创建确认消息())
            else:
                print("收到不需要的包", self.需要的消息ID, 消息)
        elif 消息.消息类型 == UDP消息类型.确认收到:
            if self.待确认消息 is not None and self.消息确认futre is not None:
                if 消息.消息ID == self.待确认消息.消息ID:
                    self.消息确认futre.set_result(True)
                    if not self.发送消息队列.empty():
                        (待发送消息, future) = self.发送消息队列.get_nowait()
                        if not self.发送无需确认是否收到消息(待发送消息):
                            raise UDP链接已关闭()
                        self.待确认消息 = 待发送消息
                        self.消息确认futre = future
                    else:
                        self.待确认消息 = None
                        self.消息确认futre = None
                else:
                    print("收到不需要的包", 消息)
        elif 消息.消息类型 == UDP消息类型.关闭:
            self.关闭()
        elif 消息.消息类型 == UDP消息类型.请求建立链接:
            await self.发送链接成功()

    async def 读一条消息(self, 超时时间: float | None = None):
        self.运行状态检测()
        if 超时时间 is not None:
            开始时间 = time.time()
            while time.time() - 开始时间 < 超时时间:
                if not self.收到的消息队列.empty():
                    return self.收到的消息队列.get_nowait()
                await asyncio.sleep(0.1)
            raise 超时()
        return await self.收到的消息队列.get()

    async def 消息重发任务(self):
        while not self.协议.链接关闭 and not self.会话关闭:
            if self.待确认消息 is not None:
                当前时间 = time.time()
                if 当前时间 - self.待确认消息.时间辍 > UDP服务器信息.消息重传时间:
                    if self.待确认消息.重传次数 < UDP服务器信息.消息重传最大次数:
                        self.协议.write(self.待确认消息(), self.地址)
                        self.待确认消息.时间辍 = 当前时间
                        self.待确认消息.重传次数 += 1
                    else:
                        self.关闭()
                        if self.消息确认futre is not None:
                            self.消息确认futre.cancel()
                        return
            await asyncio.sleep(0.1)

    async def 发送需要确认收到的消息(self, 消息: UDP消息):
        消息确认futre = self.loop.create_future()
        if self.待确认消息 is None:
            if not self.协议.write(消息(), self.地址):
                raise UDP链接已关闭()
            self.待确认消息 = 消息
            self.消息确认futre = 消息确认futre
        else:
            self.发送消息队列.put_nowait((消息, 消息确认futre))
        await 消息确认futre
        self.心跳.刷新发送()

    def 发送无需确认是否收到消息(self, 消息: UDP消息):
        if not self.协议.write(消息(), self.地址):
            raise UDP链接已关闭()
        self.心跳.刷新发送()

    def 构建数据消息(self, 数据: bytes):
        return UDP消息(UDP消息类型.数据, self.会话ID, 数据, self.消息ID生成器(), self.地址)

    async def 发送数据(self, 数据: bytes):
        消息 = self.构建数据消息(数据)
        return await self.发送需要确认收到的消息(消息)

    async def 发送链接请求(self):
        self.运行状态检测()
        await self.发送需要确认收到的消息(UDP消息(UDP消息类型.请求建立链接, self.会话ID))

    async def 发送链接成功(self):
        self.运行状态检测()
        await self.发送需要确认收到的消息(UDP消息(UDP消息类型.链接成功, self.会话ID, 地址=self.地址))
